from torch import nn
import torch
from torchvision.models import mobilenet_v3_small,resnet18

class Models(torch.nn.Module):
    def __init__(self,num_class=5):
        super(Models, self).__init__()
        self.Conv = torch.nn.Sequential(
            torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(16),
            torch.nn.ReLU(),
            torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),

            torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            torch.nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((1,1)))

        self.Classes = torch.nn.Sequential(
            torch.nn.Linear(32 * 1 * 1, 16),
            torch.nn.ReLU(),
            torch.nn.Dropout(p=0.5),
            torch.nn.Linear(16, num_class))

    def forward(self, inputs):
        x = self.Conv(inputs)
        x = x.view(-1, 32 * 1 * 1)
        x = self.Classes(x)
        return x



if __name__ == "__main__":
    model = Models()
    dummy_input = torch.ones([1,3,256,256])
    output = model(dummy_input)
    print(output.shape)
    